import sys
sys.path.insert(0, "/home/ninetripod/catkin_ws/src/uav/scripts")
sys.path.insert(0, "/home/ninetripod/catkin_ws/src/uav/scripts/visual")
# print(sys.path)
import rospy
from nav_msgs.msg import Odometry
from std_msgs.msg import String
import cv2, json, traceback, time
import numpy as np
import vtool
from peripherals import Dial, get_LED1

led1 = get_LED1()
led1.write(True)

class TaskIdSubscriber:
    def __init__(self) -> None:
        rospy.Subscriber("uav/mode", String, self.onData)
        self._data = None
        self._state = 0

    def onData(self, data):
        self._data = data.data
    
    @property
    def state(self):
        if self._state == 0:
            if self._data == '2':
                self._state = 1
        elif self._state == 1:
            if self._data == '1':
                self._state = 0
            elif self._data == '3':
                self._state = 2            
        return self._state

class Detect():

    def __init__(self,video_dev = '/dev/video17'):

        cap = cv2.VideoCapture(video_dev)
        self.cap = cap
        cap.set(cv2.CAP_PROP_FOURCC,cv2.VideoWriter_fourcc(*"MJPG"))
        cap.set(cv2.CAP_PROP_FPS, 30)

        fourcc = int(cap.get(cv2.CAP_PROP_FOURCC))
        fps = cap.get(cv2.CAP_PROP_FPS)
        width, height = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))

       
        self.pub = rospy.Publisher("uav/pixel_pos", String, queue_size = 10)
        self.Id_sub = TaskIdSubscriber()

        mtx = np.array([[389.1188600763897, 0.0, 340.332634484495], [0.0, 390.66690938053046, 301.5605063606293], [0.0, 0.0, 1.0]])
        dist = np.array([[-0.06710032864662062, -0.019528809318719448, -0.0005139028515977271, 0.0005799139558447469, -0.017234191187122435]])
        w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        self.new_camera_mtx,self.roi = cv2.getOptimalNewCameraMatrix(mtx,dist,(w,h),1,(w,h))
        self.mtx = mtx
        self.dist = dist

        x, y, w, h = self.roi
        self.video_out = cv2.VideoWriter('/home/ninetripod/Videos/video_{}.mp4'.format(time.gmtime().tm_min), fourcc, fps, (int(w), int(h)))


        # self.rate = rospy.Rate(30)

        # self.aim = self._get_aim()
        self.aim = "b0"

        self.got_time = 0
        self.start_time = time.time()

        self.record = True
        self.enable = True

    def _get_aim(self):
        dial = Dial()
        val = dial.read()

        color = None

        if val & 0b1000 == 0b1000:
            color = 'r'
        
        else:
            color = 'b'

        edges = val & 0b110

        if edges == 0b000:
            color += "0"

        elif edges == 0b010:
            color += "3"

        elif edges == 0b100:
            color += "4"

        else:
            color += "0"

        rospy.loginfo(f"set aim: {color}")
        dial.close()

        return color


    def undistort_and_cut(self,img):
        img = cv2.undistort(img, self.mtx, self.dist, None, self.new_camera_mtx)
        x,y,w,h = self.roi
        return img[y:y+h,x:x+w]

    def run_detect(self):
        while self.enable and not rospy.is_shutdown():
            try:
                ref, img = self.cap.read()

                if ref:
                    img = self.undistort_and_cut(img)

                    task_state = self.Id_sub.state
                    if task_state == 0:
                        detect_aim = vtool.decide_target(img)
                        if not detect_aim == 'None':
                            self.aim = detect_aim
                            self.pub.publish(detect_aim)
                        
                    elif task_state == 1:
                        pass

                    else:
                        img, pos = vtool.find_target(img,self.aim)
                        if not pos[0] == -1:

                            if self.got_time == 0:
                                self.got_time = time.time()

                            result = json.dumps(pos)

                            self.pub.publish(result)


                        if self.got_time != 0 and time.time() - self.got_time > 25:
                            self.record = False 
                            self.release_video_out()


                        if time.time() - self.start_time > 120: # 超时自动退出
                            self.enable = False
                            self.record = False 
                            break


                        if self.record: # 保存视频
                            self.video_out.write(img)

                else:
                    print('\r[Visual]: No frame got.',end='', flush=True)

                    if not self.got_time == 0:
                        self.enable = False
                        # self.release_cap()
                        break

            except Exception:
                traceback.print_exc()
                
                print("release camera and video")

            # self.rate.sleep()
    def release_cap(self):
        if self.cap is not None:
            self.cap.release()
            self.cap = None
    
    def release_video_out(self):
        global led1

        if self.video_out is not None:
            self.video_out.release()
            self.video_out = None
        
        if led1 is not None:
            led1.write(False)
            led1 = None
    
    def release(self):
        self.release_cap()
        self.release_video_out()

    def stop_detect(self):
        self.enable = False

if __name__ == '__main__':
    try:
        rospy.init_node("Visual_Detector")
        visual_detect = Detect()

        print("Visual_Detector init done")

        visual_detect.run_detect()


    except:
        traceback.print_exc()

    finally:
        visual_detect.release()
        print("release camera and video")



